# -*- coding: utf-8 -*-
from abc import ABCMeta, abstractmethod
from copy import deepcopy
import logging
import numpy as np
from numpy.linalg import inv, norm
import trimesh
import trimesh.transformations as tf
from autolab_core import RigidTransform
from .grasp import Grasp


class PointSuction(Grasp):
    """ Abstract grasp class for grasps with a point contact model.

    Attributes
    ----------
    configuration : :obj:`dict`
    frame : :obj:`str`
        string name of grasp reference frame (defaults to obj)
    """
    __metaclass__ = ABCMeta

    @abstractmethod
    def create_line_of_action(g, axis, width, obj, num_samples):
        """ Creates a line of action, or the points in space that the grasp traces out, from a point g in world coordinates on an object.

        Returns
        -------
        bool
            whether or not successful
        :obj:`list` of :obj:`numpy.ndarray`
            points in 3D space along the line of action
        """
        pass


class Suction3D(PointSuction):
    """ Parallel Jaw point grasps in 3D space.
    """

    def __init__(self, configuration, frame='object', grasp_id=None):
        # get parameters from configuration dict
        suction_center, suction_approach, radius, distance_to_com, score, contacts, sampled_contacts_num = \
            Suction3D.params_from_configuration(configuration)

        self.center_ = suction_center
        self.axis_ = suction_approach / np.linalg.norm(suction_approach)
        self.radius_ = radius
        self.distance_to_com_ = distance_to_com
        self.score_ = score
        self.contacts_ = contacts
        self.sampled_contacts_num_ = sampled_contacts_num
        self.frame_ = frame
        self.grasp_id_ = grasp_id

    @property
    def center(self):
        """ :obj:`numpy.ndarray` : 3-vector specifying the center of the jaws """
        return self.center_

    @center.setter
    def center(self, x):
        self.center_ = x

    @property
    def axis(self):
        """ :obj:`numpy.ndarray` : normalized 3-vector specifying the line between the jaws """
        return self.axis_

    @property
    def radius(self):
        """ float : maximum opening width of the jaws """
        return self.radius_

    @property
    def distance_to_com(self):
        """ float : maximum opening width of the jaws """
        return self.distance_to_com_

    @property
    def score(self):
        """ float : minimum opening width of the jaws """
        return self.score_

    @property
    def contacts(self):
        """ float : minimum opening width of the jaws """
        return self.contacts_

    @property
    def sampled_contacts_num(self):
        """ float : approach angle of the grasp """
        return self.sampled_contacts_num_

    @property
    def configuration(self):
        return Suction3D.configuration_from_params(self.center_, self.axis_, self.radius_, self.distance_to_com_,
                                                   self.score_, self.contacts_, self.sampled_contacts_num_)

    @property
    def frame(self):
        """ :obj:`str` : name of grasp reference frame """
        return self.frame_

    @property
    def id(self):
        """ int : id of grasp """
        return self.grasp_id_

    @frame.setter
    def frame(self, f):
        self.frame_ = f

    @staticmethod
    def distance(g1, g2, alpha=0.06):
        """ Evaluates the distance between two grasps.

        Parameters
        ----------
        g1 : :obj:`ParallelJawPtGrasp3D`
            the first grasp to use
        g2 : :obj:`ParallelJawPtGrasp3D`
            the second grasp to use
        alpha : float
            parameter weighting rotational versus spatial distance

        Returns
        -------
        float
            distance between grasps g1 and g
        """
        center_dist = np.linalg.norm(g1.center - g2.center)
        value = np.abs(g1.axis.dot(g2.axis))
        axis_dist = (1.0 / np.pi) * np.arccos(value)
        return center_dist + alpha * axis_dist

    @staticmethod
    def configuration_from_params(suction_point, suction_approach, radius, distance_to_com, quality, contact_points,
                                  sampled_points_num):
        """ Converts grasp parameters to a configuration vector. """
        if np.abs(np.linalg.norm(suction_approach) - 1.0) > 1e-5:
            raise ValueError('Illegal grasp axis. Must be norm one')
        # initialize a empty configuration dict
        configuration = {'center': suction_point, 'axis': suction_approach, 'radius': radius}
        assert len(contact_points) == sampled_points_num
        configuration['distance_to_com'] = distance_to_com
        configuration['score'] = quality
        configuration['contacts'] = contact_points
        configuration['sampled_points_num'] = sampled_points_num

        return configuration

    @staticmethod
    def params_from_configuration(configuration):
        """ Converts configuration vector into grasp parameters.

        Returns
        -------
        suction_center : :obj:`numpy.ndarray`
            center of grasp in 3D space
        grasp_axis : :obj:`numpy.ndarray`
            normalized axis of grasp in 3D space
        max_width : float
            maximum opening width of jaws
        angle : float
            approach angle
        jaw_width : float
            width of jaws
        min_width : float
            minimum closing width of jaws
        """
        if not isinstance(configuration, dict):
            raise ValueError('Configuration must be dict')
        if np.abs(np.linalg.norm(configuration['axis']) - 1.0) > 1e-5:
            raise ValueError('Illegal grasp axis. Must be norm one')
        return configuration['center'], configuration['axis'], configuration['radius'],  \
               configuration['distance_to_com'], configuration['score'], configuration['contacts'], \
               configuration['sampled_points_num']

    @property
    def unrotated_full_axis(self):
        """ Rotation matrix from canonical grasp reference frame to object reference frame. X axis points out of the
        gripper palm along the 0-degree approach direction, Y axis points between the jaws, and the Z axs is orthogonal.

        Returns
        -------
        :obj:`numpy.ndarray`
            rotation matrix of grasp
        """
        grasp_axis_y = -self.axis
        grasp_axis_x = np.array([grasp_axis_y[1], -grasp_axis_y[0], 0])
        if np.linalg.norm(grasp_axis_x) == 0:
            grasp_axis_x = np.array([1, 0, 0])
        grasp_axis_x = grasp_axis_x / norm(grasp_axis_x)
        grasp_axis_z = np.cross(grasp_axis_x, grasp_axis_y)

        R = np.c_[np.c_[grasp_axis_z, grasp_axis_x], grasp_axis_y]
        return R

    @property
    def T_grasp_obj(self):
        """ Rigid transformation from grasp frame to object frame.
        Rotation matrix is X-axis along approach direction, Y axis pointing between the jaws, and Z-axis orthogonal.
        Translation vector is the grasp center.

        Returns
        -------
        :obj:`RigidTransform`
            transformation from grasp to object coordinates
        """
        T_grasp_obj = RigidTransform(self.unrotated_full_axis, self.center, from_frame='grasp', to_frame='obj')
        return T_grasp_obj

    def close_fingers(self, obj, gripper, collision_manager, R, vis=False, check_approach=True):
        collision_free = True
        if check_approach:
            unified_contact_points = np.dot(np.append(self.contacts, self.center[np.newaxis, ...], axis=0), R[:3, :3].T)
            z_max_diff = unified_contact_points[:, 2].max() - unified_contact_points[:, 2].min()
            # print('z_max_diff: ', z_max_diff)
            if z_max_diff < gripper.z_tolerance:
                gripper_mesh = deepcopy(gripper.mesh)
                transformation = self.T_grasp_obj.matrix
                gripper_mesh.apply_transform(transformation)
                # approach_distance = (gripper.z_tolerance - 0.0005) - \
                #                     (np.dot(self.center, R[:3, :3].T)[2] - unified_contact_points[:, 2].min())
                # if approach_distance < 0:
                #     approach_distance = 0
                #
                # # print('approach_distance is : ', approach_distance)
                # translation = tf.translation_matrix(self.axis * approach_distance)
                # gripper_mesh.apply_transform(translation)
                collision_free = not collision_manager.in_collision_single(gripper_mesh, return_names=False,
                                                                           return_data=False)
            else:
                collision_free = False
            vis = False
            if vis and z_max_diff < gripper.z_tolerance:
                pc_1 = trimesh.PointCloud(self.contacts, colors=[0, 255, 0])
                # pc_2 = trimesh.PointCloud(unified_contact_points+np.array([0, 0, 0.05]), colors=[0, 0, 255])
                scene = trimesh.Scene([obj.mesh, gripper_mesh, pc_1])

                scene.show()
        # print('collision free:', collision_free)ls

        return collision_free

    @staticmethod
    def find_contact(obj, suction_point, direction, circle_sampled_points_num, radius, vis=False):
        vertices = trimesh.creation.cylinder(radius=radius, height=2, sections=circle_sampled_points_num,
                                             transform=trimesh.transformations.translation_matrix([0, 0, 1])).vertices
        predefined_vertices = vertices[vertices[:, 2] == 0]
        R = tf.rotation_matrix(tf.angle_between_vectors(direction, np.array([0, 0, 1])),
                               tf.vector_product(direction, np.array([0, 0, 1])))
        circle_vertices = np.dot(predefined_vertices, R[:3, :3]) + suction_point
        directions = np.array([direction] * (circle_sampled_points_num + 1))
        locations, index_ray, index_tri = obj.mesh.ray.intersects_location(circle_vertices, directions,
                                                                           multiple_hits=False)
        # vis = True
        if vis:
            pc_1 = trimesh.PointCloud(circle_vertices, colors=[0, 255, 0])
            if len(locations) > 0:
                pc_2 = trimesh.PointCloud(locations, colors=[0, 0, 255])
                scene = trimesh.Scene([obj.mesh, pc_1, pc_2])
            else:
                scene = trimesh.Scene([obj.mesh, pc_1])
            scene.show()

        return locations, index_ray, index_tri, R

    @staticmethod
    def suction_from_point_and_axis_on_mesh(obj, suction_candidate, suction_approach, radius=0.01, backup=0.1,
                                            circle_sampled_points_num=36, vis=False):
        """
        """
        # transform to grid basis
        suction_candidate_world = suction_candidate.point
        suction_approach = suction_approach / np.linalg.norm(suction_approach)
        # backup point along the normal
        point_backup = deepcopy(suction_candidate_world)
        point_backup -= backup * suction_approach

        contacts, ray_index, _, R = Suction3D.find_contact(obj, suction_point=point_backup,
                                                           direction=suction_approach,
                                                           circle_sampled_points_num=circle_sampled_points_num,
                                                           radius=radius, vis=vis)

        # contacts points num should equal to circle_sampled_points+1
        if len(contacts) == (circle_sampled_points_num + 1):
            ray_index_unique = np.unique(ray_index)
            if len(ray_index_unique) == (circle_sampled_points_num + 1):
                contacts_found = True
            else:
                contacts_found = False
        else:
            contacts_found = False

        grasp = None
        if contacts_found:
            contact_ray_tuple = [(contact_, index_) for contact_, index_ in zip(contacts, ray_index)]
            sorted_contacts = sorted(contact_ray_tuple, key=lambda s: s[1])
            distance_ = suction_candidate_world - sorted_contacts[0][0]
            if np.linalg.norm(distance_) > 0.001:
                contacts_found = False
            else:
                curvature_satisfy = True
                suction_point = sorted_contacts[0][0]
                contact_points = np.array([contact[0] for contact in sorted_contacts[1:]])
                closest_vertices = trimesh.proximity.ProximityQuery(obj.mesh).vertex(contact_points)[1]
                curvature = obj.mean_curvature[closest_vertices]
                min_curvature = curvature.min()
                max_curvature = curvature.max()
                if max_curvature > 60 or min_curvature < -20:
                    curvature_satisfy = False
                    contacts_found = False

                # from m to mm
                # fit a plane C is point on plane, N is normal of the point
                if curvature_satisfy:
                    C, N = trimesh.points.plane_fit(contact_points * 1000)
                    # angle between N and suction_approach
                    discount_rate = np.dot(N, suction_approach)
                    if discount_rate < 0:
                        discount_rate = -discount_rate
                    com = obj.mesh.center_mass
                    distance_to_com = np.linalg.norm(np.cross(suction_approach, point_backup - com)) / np.linalg.norm(
                        suction_approach)

                    distances = trimesh.points.point_plane_distance(contact_points * 1000, plane_normal=N, plane_origin=C)
                    distances_norm = np.linalg.norm(distances)
                    distances_var = distances_norm / np.sqrt(circle_sampled_points_num)
                    quality = np.exp(-distances_var) * discount_rate
                    configuration = Suction3D.configuration_from_params(suction_point, suction_approach, radius,
                                                                        distance_to_com,
                                                                        quality, contact_points,
                                                                        circle_sampled_points_num)
                    grasp = Suction3D(configuration)

        if not contacts_found:
            logging.debug('No sufficient contacts found for suction')

        return contacts_found, grasp, R  # relative to object
